import torch.nn as nn

class dualAE(nn.Module):
	def __init__(self, in_sz, hd_sz):
		super(dualAE, self).__init__()

		base = hd_sz
		self.encoder1 = nn.Sequential(
			nn.Linear(in_sz, base * 2),
			nn.ReLU(),
			nn.BatchNorm1d(base * 2),
			nn.Linear(base * 2, base),
			nn.ReLU(),
			)
		self.encoder2 = nn.Sequential(
			nn.Linear(in_sz, base * 2),
			nn.ReLU(),
			nn.BatchNorm1d(base * 2),
			nn.Linear(base * 2, base),
			nn.ReLU(),
			)
		self.decoder1 = nn.Sequential(
			nn.BatchNorm1d(base),
			nn.Linear(base, base * 2),
			nn.ReLU(),
			nn.BatchNorm1d(base * 2),
			nn.Linear(base * 2, in_sz),
			nn.ReLU()
			)
		self.decoder2 = nn.Sequential(
			nn.BatchNorm1d(base),
			nn.Linear(base, base * 2),
			nn.ReLU(),
			nn.BatchNorm1d(base * 2),
			nn.Linear(base * 2, in_sz),
			nn.ReLU()
			)

	def forward(self, x, channel):
		if channel == 1:
			return self.decoder1(self.encoder1(x)), self.encoder1(x)
		else:
			return self.decoder2(self.encoder2(x)), self.encoder2(x)
			
